"""
作用：读入一个zarr数组，然后对其进行有规则地打乱
"""
#!/usr/bin/env python3
import argparse
import zarr

def main():

    parser = argparse.ArgumentParser(description="Shuffle zarr dataset with options.")
    parser.add_argument('-f', '--from_path', type=str, required=True, help='Path to the input zarr file/folder to shuffle')
    parser.add_argument('-t', '--to_path', type=str, default=None, help='Path to save the shuffled zarr file/folder (default: same folder as from_path)')
    parser.add_argument('-b', '--batch_size', type=int, default=32, help='Batch size')
    parser.add_argument('-H', '--horizon', type=int, default=1, help='Horizon')
    parser.add_argument('-n', type=int, default=3, help='Number of groups/patterns/robot-arms to shuffle')
    parser.add_argument('-a', '--array', type=str, default=None, help='Array name under group (if input is a group)')
    args = parser.parse_args()

    from_path = args.from_path
    to_path = args.to_path
    if to_path is None:
        import os
        base, name = os.path.split(from_path.rstrip('/'))
        to_path = os.path.join(base, f"{name}-shuffled")

    import numpy as np
    import zarr
    import sys

    # Open source zarr
    src_zarr = zarr.open(from_path, mode='r')
    # If src_zarr is a zarr.Array, use directly; if zarr.Group, pick first array
    if isinstance(src_zarr, zarr.Array):
        arr = src_zarr
    elif isinstance(src_zarr, zarr.Group):
        if args.array is not None:
            # Support nested path like 'data/action'
            try:
                arr = src_zarr
                for part in args.array.strip('/').split('/'):
                    arr = arr[part]
                if not isinstance(arr, zarr.Array):
                    print(f"Error: Path '{args.array}' does not point to a zarr Array!", file=sys.stderr)
                    sys.exit(1)
            except Exception as e:
                print(f"Error: Array path '{args.array}' not found in group! Exception: {e}", file=sys.stderr)
                sys.exit(1)
        else:
            arr_keys = [k for k in src_zarr.keys() if isinstance(src_zarr[k], zarr.Array)]
            if not arr_keys:
                print("Error: No array found in the zarr file!", file=sys.stderr)
                sys.exit(1)
            arr = src_zarr[arr_keys[0]]
    else:
        print("Error: Input is neither a zarr Array nor Group!", file=sys.stderr)
        sys.exit(1)

    # Ensure N is a plain int (sometimes numpy returns np.int32 etc)
    N = arr.shape[0]
    import numpy as np
    if isinstance(N, (int, float)):
        N = int(N)
    elif isinstance(N, (np.integer, np.floating)):
        N = int(N)
    else:
        print(f"Error: arr.shape[0] is of unsupported type {type(N)}!", file=sys.stderr)
        sys.exit(1)
    n = args.n
    batch_size = args.batch_size

    # 1. Check divisibility
    if N % n != 0:
        print(f"Error: N={N} is not divisible by n={n}!", file=sys.stderr)
        sys.exit(1)
    seg_len = N // n
    if seg_len % batch_size != 0:
        print(f"Error: Each segment length {seg_len} is not divisible by batch_size={batch_size}!", file=sys.stderr)
        sys.exit(1)

    # 2. Split and shuffle each segment
    indices = np.arange(N)
    segments = []
    for i in range(n):
        start = i * seg_len
        end = (i + 1) * seg_len
        seg_idx = indices[start:end].copy()
        np.random.shuffle(seg_idx)
        segments.append(seg_idx)

    # 3. Interleave by batch_size
    num_batches = seg_len // batch_size
    out_indices = []
    for b in range(num_batches):
        for seg in segments:
            out_indices.extend(seg[b*batch_size:(b+1)*batch_size])

    # 4. Write to new zarr (do not overwrite original)
    import os
    if os.path.exists(to_path):
        print(f"Error: Output path {to_path} already exists!", file=sys.stderr)
        sys.exit(1)
    store = zarr.DirectoryStore(to_path)
    zarr_out = zarr.group(store=store)
    # Prepare dataset creation kwargs safely
    ds_kwargs = {
        'shape': arr.shape,
        'dtype': arr.dtype
    }
    # Only add optional attributes if present
    if hasattr(arr, 'chunks') and arr.chunks is not None:
        ds_kwargs['chunks'] = arr.chunks
    if hasattr(arr, 'compressor') and arr.compressor is not None:
        ds_kwargs['compressor'] = arr.compressor
    if hasattr(arr, 'fill_value') and arr.fill_value is not None:
        ds_kwargs['fill_value'] = arr.fill_value
    # Use a safe name
    arr_name = getattr(arr, 'name', 'data')
    # Fix: if arr_name is empty, root, or None, use 'data' to avoid group conflict
    if not arr_name or arr_name == '/' or arr_name == '.':
        arr_name = 'data'
    arr_out = zarr_out.create_dataset(arr_name, **ds_kwargs)
    # Copy data in new order
    # zarr does not support advanced numpy indexing, so assign row by row
    out_indices = np.array(out_indices, dtype=int)
    for i, idx in enumerate(out_indices):
        arr_out[i] = arr[idx]
    print(f"Shuffle done! Output saved to {to_path}")

if __name__ == "__main__":
    main()